{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Verify S3_BUCKET Bucket Creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import time\n",
    "\n",
    "session = boto3.session.Session()\n",
    "region = session.region_name\n",
    "account_id = boto3.client(\"sts\").get_caller_identity().get(\"Account\")\n",
    "\n",
    "s3 = boto3.Session().client(service_name=\"s3\", region_name=region)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "setup_s3_bucket_passed = False\n",
    "%store -r airflow_bucket_name\n",
    "%store -r s3_mwaa_private_path\n",
    "%store -r s3_mwaa_dags_private_path\n",
    "%store -r airflow_env_name\n",
    "%store -r airflow_vpc_name\n",
    "%store -r team_role_arn\n",
    "%store -r airflow_sg_id\n",
    "%store -r airflow_subnet_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!aws s3 ls $s3_mwaa_private_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from botocore.client import ClientError\n",
    "\n",
    "response = None\n",
    "\n",
    "try:\n",
    "    response = s3.head_bucket(Bucket=airflow_bucket_name)\n",
    "    print(response)\n",
    "    setup_s3_bucket_passed = True\n",
    "except ClientError as e:\n",
    "    print(\"[ERROR] Cannot find bucket {} in {} due to {}.\".format(airflow_bucket_name, response, e))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %store setup_s3_bucket_passed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Managed Apache Airflow Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mwaa = boto3.client(\"mwaa\")\n",
    "\n",
    "s3_mwaa_bucket_arn = \"arn:aws:s3:::{}\".format(airflow_bucket_name)\n",
    "\n",
    "airflow_env_arn = mwaa.create_environment(\n",
    "    DagS3Path=\"dags\",\n",
    "    ExecutionRoleArn=team_role_arn,\n",
    "    AirflowVersion=\"1.10.12\",\n",
    "    WebserverAccessMode=\"PUBLIC_ONLY\",\n",
    "    LoggingConfiguration={\n",
    "        \"DagProcessingLogs\": {\"Enabled\": True, \"LogLevel\": \"ERROR\"},\n",
    "        \"SchedulerLogs\": {\"Enabled\": True, \"LogLevel\": \"ERROR\"},\n",
    "        \"TaskLogs\": {\"Enabled\": True, \"LogLevel\": \"INFO\"},\n",
    "        \"WebserverLogs\": {\"Enabled\": True, \"LogLevel\": \"ERROR\"},\n",
    "        \"WorkerLogs\": {\"Enabled\": True, \"LogLevel\": \"ERROR\"},\n",
    "    },\n",
    "    MaxWorkers=3,\n",
    "    Name=airflow_env_name,\n",
    "    NetworkConfiguration={\n",
    "        \"SecurityGroupIds\": [\n",
    "            airflow_sg_id,\n",
    "        ],\n",
    "        \"SubnetIds\": airflow_subnet_ids,\n",
    "    },\n",
    "    RequirementsS3ObjectVersion=\"latest\",\n",
    "    RequirementsS3Path=\"requirements.txt\",\n",
    "    SourceBucketArn=s3_mwaa_bucket_arn,\n",
    "    EnvironmentClass=\"mw1.small\",\n",
    ")\n",
    "\n",
    "%store airflow_env_arn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Please be patient this can take around 15 Minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_airflow_check():\n",
    "    response = mwaa.get_environment(Name=airflow_env_name)\n",
    "    mwaa_status = response[\"Environment\"][\"Status\"]\n",
    "    return mwaa_status\n",
    "\n",
    "\n",
    "mwaa_status = \"CREATING\"\n",
    "\n",
    "print(\"Checking to see if MWAA Env: {} is ready.\".format(airflow_env_name))\n",
    "\n",
    "while get_airflow_check() != \"AVAILABLE\":\n",
    "    mwaa_status\n",
    "    time.sleep(60)\n",
    "    print(\"Still waiting for MWAA Environment...\")\n",
    "\n",
    "print(\"Sucess! MWAA Env: {} is ready!\".format(airflow_env_name))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PLEASE MAKE SURE THAT THE ABOVE COMMAND RAN SUCESSFULLY BEFORE CONTINUING"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = mwaa.create_web_login_token(\n",
    "    Name=airflow_env_name\n",
    ")\n",
    "\n",
    "webServerHostName = response[\"WebServerHostname\"]\n",
    "webToken = response[\"WebToken\"]\n",
    "airflowUIUrl = 'https://'+webServerHostName+'/aws_mwaa/aws-console-sso?login=true#'+webToken\n",
    "print(\"Here is your AirflowUI Url:\"\\n)\n",
    "airflowUIUrl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Release Resources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%html\n",
    "\n",
    "<p><b>Shutting down your kernel for this notebook to release resources.</b></p>\n",
    "<button class=\"sm-command-button\" data-commandlinker-command=\"kernelmenu:shutdown\" style=\"display:none;\">Shutdown Kernel</button>\n",
    "        \n",
    "<script>\n",
    "try {\n",
    "    els = document.getElementsByClassName(\"sm-command-button\");\n",
    "    els[0].click();\n",
    "}\n",
    "catch(err) {\n",
    "    // NoOp\n",
    "}    \n",
    "</script>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%javascript\n",
    "\n",
    "try {\n",
    "    Jupyter.notebook.save_checkpoint();\n",
    "    Jupyter.notebook.session.delete();\n",
    "}\n",
    "catch(err) {\n",
    "    // NoOp\n",
    "}"
   ]
  }
 ],
 "metadata": {
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "Python 3 (Data Science)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-east-1:081325390199:image/datascience-1.0"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
